package fem2.strategies;

import inf.jlinalg.BLAM;
import inf.jlinalg.IVector;
import fem2.Model;

/**
 * Arc length solution strategy with sphere constraint
 * 
 * @author hbui
 * 
 */
public class ArcLengthSolutionStrategyWithSphereConstraint extends ArcLengthSolutionStrategy {

	private double s;
	private double[] sArray = null;
	private double psi;
	private int[] dofsIndex = null;

	private int step = 0;

	/**
	 * 
	 * @param m
	 *            model
	 * @param s
	 *            stepping parameter
	 * @param psi
	 *            scaling factor
	 */
	public ArcLengthSolutionStrategyWithSphereConstraint(Model m, double s, double psi) {
		super(m);
		this.s = s;
		this.psi = psi;
	}

	/**
	 * 
	 * @param m
	 *            model
	 * @param s
	 *            stepping parameter
	 * @param psi
	 *            scaling factor
	 * @param dofsIndex
	 *            index of constrainted dofs
	 */
	public ArcLengthSolutionStrategyWithSphereConstraint(Model m, double s, double psi,
			int... dofsIndex) {
		super(m);
		this.s = s;
		this.psi = psi;
		this.dofsIndex = dofsIndex;
	}

	/**
	 * 
	 * @param m
	 *            model
	 * @param sArray
	 *            array of stepping parameter
	 * @param psi
	 *            scaling factor
	 */
	public ArcLengthSolutionStrategyWithSphereConstraint(Model m, double[] sArray, double psi) {
		super(m);
		this.psi = psi;
		this.sArray = sArray;
		this.s = sArray[0];
		setNumberOfTimeStep(sArray.length);
	}

	/**
	 * 
	 * @param m
	 * @param sArray
	 *            array of stepping parameter
	 * @param psi
	 * @param dofsIndex
	 *            index of constrainted dofs
	 */
	public ArcLengthSolutionStrategyWithSphereConstraint(Model m, double[] sArray, double psi,
			int... dofsIndex) {
		super(m);
		this.psi = psi;
		this.sArray = sArray;
		this.s = sArray[0];
		setNumberOfTimeStep(sArray.length);
		this.dofsIndex = dofsIndex;
	}

	@Override
	protected double getNextSteppingParameter() {
		if (sArray != null) {
			s = sArray[step++];
		}
		return s;
	}

	@Override
	protected double getCurrentSteppingParameter() {
		return s;
	}

	@Override
	protected double getScalingParameter() {
		return psi;
	}

	@Override
	protected double computeConstraintValue(IVector u, IVector u_old, double lambda,
			double lambda_old) {
		if (dofsIndex == null) {
			IVector du = (IVector) u.clone();
			BLAM.add(-1.0, u_old, du);
			return Math.sqrt(Math.pow(BLAM.norm(BLAM.NORM_TWO, du), 2)
					+ Math.pow(psi * (lambda - lambda_old), 2))
					- s;
		} else {
			double norm_du_pow_2 = 0.0;
			for (int i = 0; i < dofsIndex.length; i++) {
				norm_du_pow_2 += Math.pow(u.get(dofsIndex[i]) - u_old.get(dofsIndex[i]), 2);
			}
			return Math.sqrt(norm_du_pow_2 + Math.pow(psi * (lambda - lambda_old), 2)) - s;
		}
	}

	@Override
	protected void computeGradientOfConstraintWrtU(IVector u, IVector u_old, double lambda,
			double lambda_old, IVector dfdu) {
		if (dofsIndex == null) {
			int n = u.getSize();
			double f = computeConstraintValue(u, u_old, lambda, lambda_old) + s;
			for (int i = 0; i < n; i++) {
				dfdu.set(i, (u.get(i) - u_old.get(i)) / f);
			}
		} else {
			double f = computeConstraintValue(u, u_old, lambda, lambda_old) + s;
			for (int i = 0; i < dofsIndex.length; i++) {
				dfdu.set(i, (u.get(dofsIndex[i]) - u_old.get(dofsIndex[i])) / f);
			}
		}
	}

	@Override
	protected double computeDerivativeOfConstraintWrtLambda(IVector u, IVector u_old,
			double lambda, double lambda_old) {
		double f = computeConstraintValue(u, u_old, lambda, lambda_old) + s;
		return Math.pow(psi, 2) * (lambda - lambda_old) / f;
	}

}
